import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
import torch
from transformers import BertTokenizer, BertModel, BertForMaskedLM

# 可选：如果您想了解发生的信息，请按以下步骤logger
import logging
logging.basicConfig(level=logging.INFO)

# 加载预训练的模型标记器（词汇表）
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')


# 标记输入
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = tokenizer.tokenize(text)


# 用“BertForMaskedLM”掩盖我们试图预测的标记`
masked_index = 8
tokenized_text[masked_index] = '[MASK]'
assert tokenized_text == ['[CLS]', 'who', 'was', 'jim', 'henson', '?', '[SEP]', 'jim', '[MASK]', 'was', 'a', 'puppet', '##eer', '[SEP]']

# 将标记转换为词汇索引
indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
# 定义与第一句和第二句相关的句子A和B索引（见论文）
segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]

# 将输入转换为PyTorch张量
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])